include "stdlib.h";

trait Allocator {
    function alloc(n: Size): Ptr[Void];
    function free(ptr: Ptr[Void]): Void;

    // function safeAlloc(n: Size): Option[Ptr[Void]] {
    //     var ptr = this.alloc(n);
    //     return if ptr == null then None else Some(ptr);
    // }

    function calloc(n: Size): Ptr[Void] {
        var ptr = this.Allocator.alloc(n);
        memset(ptr, 0, n);
        return ptr;
    }
}

struct SimpleAllocator {
    public var alloc: function (Size) -> Ptr[Void];
    public var calloc: function (Size) -> Ptr[Void];
    public var free: function (Ptr[Void]) -> Void;
}

implement Allocator for SimpleAllocator {
    public function alloc(n: Size): Ptr[Void] {
        return (this.alloc)(n);
    }

    public function calloc(n: Size): Ptr[Void] {
        return (this.calloc)(n);
    }

    public function free(ptr: Ptr[Void]) {
        (this.free)(ptr);
    }
}

function callocWrapper(n: Size): Ptr[Void] {
    return calloc(n, 1);
}

// will be used as the default Ptr[Void] allocator
var mallocator: Box[Allocator] = struct SimpleAllocator {
    alloc: malloc,
    calloc: callocWrapper,
    free: free,
} as Box[Allocator];

/**
 * Preallocates a region of memory, then provides pointers to unused memory in
 * linear order.
 *
 * Freeing individual allocations is a noop; instead, call `reset()` to free
 * everything at once.
 */
struct LinearAllocator {
    var parent: Box[Allocator];
    var start: Ptr[Void];
    var curPtr: Ptr[Void];
    var capacity: Size;

    public static function new(parent: Box[Allocator], capacity: Size) {
        var start = parent.alloc(capacity);
        return struct LinearAllocator {
            parent,
            start,
            curPtr: start,
            capacity
        };
    }

    public function alloc(n: Size): Option[Ptr[Void]] {
        if this.capacity - (this.curPtr as Ptr[Uint8] - this.start as Ptr[Uint8]) as Size >= n {
            var next = this.curPtr;
            this.curPtr = this.curPtr as Ptr[Uint8] + n;
            return Some(next);
        } else {
            return None;
        }
    }

    public function reset() {
        this.curPtr = this.start;
    }

    public function remaining(): Size {
        return this.capacity - (this.curPtr as Ptr[Uint8] - this.start as Ptr[Uint8]) as Size;
    }

    public function destroy() {
        this.parent.free(this.start);
    }
}

implement Allocator for LinearAllocator {
    public function alloc(s): Ptr[Void] {
        match this.alloc(s) {
            Some(x) => return x;
            default => {
                panic("failed to allocate from LinearAllocator");
                return null;
            }
        }
    }

    public function free(p) {}
}

/**
 * Preallocates a region of memory for use as a custom stack. Values should
 * be freed in LIFO order.
 */
struct StackAllocator {
    var parent: Box[Allocator];
    var start: Ptr[Void];
    var curPtr: Ptr[Void];
    var capacity: Size;

    public static function new(parent: Box[Allocator], capacity: Size) {
        var curPtr = parent.alloc(capacity);
        return struct StackAllocator {
            parent,
            start: curPtr,
            curPtr,
            capacity
        };
    }

    public function alloc(n: Size): Option[Ptr[Void]] {
        if this.capacity >= n {
            var next = this.curPtr;
            this.curPtr = (this.curPtr as Ptr[Uint8]) + n;
            this.capacity -= n;
            return Some(next);
        } else {
            return None;
        }
    }

    public function free(ptr: Ptr[Void]) {
        this.capacity += (this.curPtr as Ptr[Uint8] - ptr as Ptr[Uint8]) as Size;
        this.curPtr = ptr;
    }

    public function remaining(): Size {
        return this.capacity;
    }

    public function destroy() {
        this.parent.free(this.start);
    }
}

implement Allocator for StackAllocator {
    public function alloc(s): Ptr[Void] {
        match this.alloc(s) {
            Some(x) => return x;
            default => {
                panic("failed to allocate from StackAllocator");
                return null;
            }
        }
    }

    public function free(p) {
        this.free(p);
    }
}

struct PoolNode[T] {
    var val: T;
    var next: Ptr[PoolNode[T]];
}

struct PoolAllocator[T] {
    var parent: Box[Allocator];
    var capacity: Size;
    var next: Ptr[PoolNode[T]];
    var start: Ptr[PoolNode[T]];

    public static function new(parent: Box[Allocator], count: Size) {
        var start = parent.alloc(count * sizeof PoolNode[T]) as Ptr[PoolNode[T]];
        var pool = struct PoolAllocator[T] {
            parent,
            start,
            capacity: count,
            next: null
        };
        for i in 0 ... count - 1 {
            pool.start[i].next = start + i + 1;
        }
        pool.start[count - 1].next = null;
        pool.next = pool.start;
        return pool;
    }

    public function alloc(n: Size): Option[Ptr[T]] {
        if this.capacity > 0 && this.next != null {
            --this.capacity;
            var next = this.next;
            this.next = next.next;
            next.next = null;
            return Some(next as Ptr[T]);
        } else {
            return None;
        }
    }

    public function free(ptr: Ptr[T]) {
        ++this.capacity;
        var recycled = ptr as Ptr[PoolNode[T]];
        recycled.next = this.next;
        this.next = recycled;
    }

    public function remaining(): Size {
        return this.capacity;
    }

    public function destroy() {
        this.parent.free(this.start);
    }
}

// implement Allocator for PoolAllocator[T] {
//     public function alloc(s): Ptr[Void] {
//         match this.alloc(s) {
//             Some(x) => return x;
//             default => {
//                 panic("failed to allocate from PoolAllocator");
//                 return null;
//             }
//         }
//     }

//     public function free(p) {
//         this.free(p);
//     }
// }

public function copy[T](source: Ptr[T], dest: Ptr[T]): Void {
    memcpy(dest, source, sizeof T);
}

public function clear[T](dest: Ptr[T]): Void {
    memset(dest, 0, sizeof T);
}

public function move[T](source: Ptr[T], dest: Ptr[T]): Void {
    copy(source, dest);
    clear(source);
}
